#include "client.h"
#include <string.h>
#include <errno.h>
#include <limits.h>
#include "log.h"
#include "fifo.h"
#include "io.h"
#include "socks.h"
#include "http.h"


const char *src_proto_str(enum src_proto p) {
    const char * strs[] = {
            "SOCKS4",
            "SOCKS4a",
            "SOCKS5",
            "HTTP",
    };
    if (p >= SP_INVALID) {
        return "?";
    }
    return strs[p];
}

struct request *request_realloc(struct request *old) {
    struct request *new;
    size_t newsize = old ? old->buf.size + REQ_CHUNK_LEN : REQ_CHUNK_LEN;

    if (!(new = realloc(old, offsetof(struct request, buf.raw) + newsize))) {
        errno = ENOMEM;
        return NULL;
    }

    if (old == NULL) {
        memset(new, 0, offsetof(struct request, buf.raw));
        new->buf.pos = new->buf.raw;
    } else
        new->buf.pos += (void *) new - (void *) old;

    new->buf.size = newsize;

    return new;
}

static void request_reset(struct request *req) {
    size_t bufsize = req->buf.size;
    memset(req, 0, offsetof(struct request, buf.raw));
    req->buf.pos = req->buf.raw;
    req->buf.size = bufsize;
}

static struct client *client_alloc(struct client_ctx *ctx) {
    struct client *c;

    if ((c = TAILQ_FIRST(&ctx->client_pool.unuse))) {
        TAILQ_REMOVE(&ctx->client_pool.unuse, c, entry);
        --ctx->client_pool.unuseTotal;
    } else if ((c = malloc(sizeof(*c)))) {
        c->lout = fifo_new(ctx->cfg->buf_size);
        c->rout = fifo_new(ctx->cfg->buf_size);
        ev_init(&c->levio, NULL);
        ev_init(&c->revio, NULL);
        ev_init(&c->evtimer, NULL);
        c->levio.data = c->revio.data = c->evtimer.data = c;
        c->req = request_realloc(NULL);
        if (!c->lout || !c->rout || !c->req) {
            fifo_del(c->lout);
            fifo_del(c->rout);
            free(c->req);
            free(c);
            c = NULL;
        }
    }

    if (c)
        TAILQ_INSERT_HEAD(&ctx->client_pool.inuse, c, entry);

    return c;
}

static void client_free(struct client *c) {
    TAILQ_REMOVE(&c->ctx->client_pool.unuse, c, entry);
    --c->ctx->client_pool.unuseTotal;
    free(c->req);
    fifo_del(c->lout);
    fifo_del(c->rout);
    free(c);
}

static void client_collection(struct client *c) {
    TAILQ_REMOVE(&c->ctx->client_pool.inuse, c, entry);
    if (c->ctx->client_pool.unuseTotal >= c->ctx->cfg->open_max >> 4)
        client_free(c);
    else {
        TAILQ_INSERT_TAIL(&c->ctx->client_pool.unuse, c, entry);
        ++c->ctx->client_pool.unuseTotal;
    }
}

static int srv_info_cmp(const void *a, const void *b)
{
    return ((const struct server_info *)a)->weight - ((const struct server_info *)b)->weight;
}

void update_server_weight(struct client *c, int n)
{
    int i;
    struct server_info *srv = (struct server_info *) c->srv;
    struct cfg *cfg = (struct cfg *) c->ctx->cfg;

    if (n > 0) {
        if (srv->weight < SERVER_WEIGHT_MAX)
            ++srv->weight;
    } else if (n < 0 && srv->weight > 0) {
        if (--srv->weight <= 0)
            LOGE("Server %s invalid", srv->addr);
    } else
        return;

    qsort(cfg->server_infos + cfg->server_available_idx, cfg->server_cnt - cfg->server_available_idx,
          sizeof(*cfg->server_infos), srv_info_cmp);

    for (i = cfg->server_available_idx; i < cfg->server_cnt; ++i) {
        if (cfg->server_infos[i].weight >= 1) {
            cfg->server_available_idx = i;
            break;
        }
    }
    if (i >= cfg->server_cnt) {
        LOGW("All the server is invalid, now we reinitialization");
        for (i = 0; i < cfg->server_cnt; ++i)
            cfg->server_infos[i].weight = cfg->server_initial_weight;
        cfg->server_available_idx = 0;
        return;
    }
    if (cfg->server_infos[i].weight > cfg->server_initial_weight) {
        int sub = cfg->server_infos[0].weight - cfg->server_initial_weight;
        for (i = 0; i < cfg->server_cnt; ++i)
            cfg->server_infos[i].weight -= sub;
    }
}

static const struct server_info *random_choice_server(struct client_ctx *ctx) {
    int i;
    int rand, total;
    struct server_info *ret = ctx->cfg->server_infos + ctx->cfg->server_cnt - 1;

    for (total = 0, i = ctx->cfg->server_available_idx; i < ctx->cfg->server_cnt; ++i)
        total += ctx->cfg->server_infos[i].weight;

    rand = rand_r(&ctx->rand_seed) % total;

    for (total = 0, i = ctx->cfg->server_available_idx; i < ctx->cfg->server_cnt; ++i) {
        total += ctx->cfg->server_infos[i].weight;
        if (rand <= total) {
            ret = ctx->cfg->server_infos + i;
            break;
        }
    }

    return ret;
}

static struct client *create(struct client_ctx *ctx) {
    struct client *c = client_alloc(ctx);

    if (c) {
        c->ctx = ctx;

        c->lsock = -1;
        c->rsock = -1;
        c->laddrinfo[0] = '?';
        c->laddrinfo[1] = '\0';
        c->raddrinfo[0] = '?';
        c->raddrinfo[1] = '\0';
        c->srv = random_choice_server(ctx);
        fifo_reset(c->lout);
        fifo_reset(c->rout);
        c->resolv = NULL;
        c->resolv_cb = NULL;
        c->crypt.flags = 0;

        c->src_proto = SP_INVALID;
        c->dst_info[0] = '?';
        c->dst_info[1] = '\0';

        request_reset(c->req);
    }

    return c;
}

void client_destroy(struct client *c) {
    resolv_cancel(c->resolv, c->resolv_cb, c);
    crypt_destroy(&c->crypt);

    ev_timer_stop(c->ctx->evloop, &c->evtimer);

    if (c->lsock >= 0) {
        ev_io_stop(c->ctx->evloop, &c->levio);
        close(c->lsock);
        c->lsock = -1;
    }

    if (c->rsock >= 0) {
        ev_io_stop(c->ctx->evloop, &c->revio);
        close(c->rsock);
        c->rsock = -1;
    }

    client_collection(c);
}

ssize_t client_write2sock(int fd, const void *data, size_t nbytes, struct client *c) {
    ssize_t n, total = 0;

    if ((n = write(fd, data, nbytes)) == (ssize_t)nbytes)
        return n;

    if (n < 0) {
        if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
            n = 0;
        else {
            client_log_socket_error(LOG_WARNING, c, fd, "send error");
            return n;
        }
    }

    total += n;

    struct fifo *outfifo;
    int peerfd;
    ev_io *io, *peerio;
    if (fd == c->lsock) {
        outfifo = c->lout;
        io = &c->levio;
        peerfd = c->rsock;
        peerio = &c->revio;
    } else {
        outfifo = c->rout;
        io = &c->revio;
        peerfd = c->lsock;
        peerio = &c->levio;
    }
    n = fifo_write(outfifo, data + total, nbytes - total);
    ev_io_stop(c->ctx->evloop, io);
    ev_io_set(io, fd, EV_WRITE);
    ev_io_start(c->ctx->evloop, io);
    if (peerfd >= 0)
        ev_io_stop(c->ctx->evloop, peerio);
    total += n;

    return total;
}

/* log the socket error use errno */
void client_log_socket_error(int loglevel, const struct client *c, int fd,
        const char *dowhat) {
    if (fd == c->lsock)
        LOG(loglevel, "%s %s: %s", c->laddrinfo, dowhat, strerror(errno));
    else
        LOG(loglevel, "%s(%s) %s: %s", c->srv->addr, c->raddrinfo, dowhat, strerror(errno));
}

/* return -1 if client destroyed */
int client_write_event(EV_P_ ev_io *w, struct client *c) {
    struct fifo *fifo = w->fd == c->rsock ? c->rout : c->lout;
    if (fifo_read_to_fd(fifo, w->fd) < 0) {
        client_log_socket_error(LOG_WARNING, c, w->fd, "send error");
        goto DESTROY;
    }

    if (fifo_is_empty(fifo)) {
        ev_io_stop(EV_A_ w);
        ev_io_set(w, w->fd, EV_READ);
        ev_io_start(EV_A_ w);
        int peerfd = w->fd == c->lsock ? c->rsock : c->lsock;
        if (peerfd >= 0) {
            ev_io *peerio = w->fd == c->lsock ? &c->revio : &c->levio;
            ev_io_start(EV_A_ peerio);
        }
    }

    return 0;

DESTROY:
    client_destroy(c);
    return -1;
}

/* return -1 if client destroyed */
static int relay_read_event(int fd, struct client *c) {
    ssize_t n;
    client_update_activity(c);

    if ((n = read(fd, c->ctx->inbuf, c->ctx->inbuflen)) < 0) {
        if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
            return 0;
        client_log_socket_error(LOG_WARNING, c, fd, "recv error");
        goto DESTROY;
    }
    if (!n)
        goto DESTROY;

    if (fd == c->lsock) {
        if ((n = crypt_en(&c->crypt, c->ctx->outbuf, c->ctx->outbuflen,
                c->ctx->inbuf, n)) < 0) {
            LOGW("Encrypt error");
            goto DESTROY;
        }
        if (client_write2sock(c->rsock, c->ctx->outbuf, n, c) < 0)
            goto DESTROY;
    } else {
        if ((n = crypt_de(&c->crypt, c->ctx->outbuf, c->ctx->outbuflen,
                c->ctx->inbuf, n)) < 0) {
            LOGW("Decrypt error");
            goto DESTROY;
        }
        if (client_write2sock(c->lsock, c->ctx->outbuf, n, c) < 0)
            goto DESTROY;
    }

    return 0;

DESTROY:
    client_destroy(c);
    return -1;
}

void relay_handler(EV_P_ ev_io *w, int revents) {
    struct client *c = w->data;

    if ((revents & EV_READ) && relay_read_event(w->fd, c))
        return;

    if (revents & EV_WRITE)
        client_write_event(EV_A_ w, c);
}

static void hello_handler(EV_P_ ev_io *w, int revents) {
    ssize_t n;
    EV_CB_DECLARE(ev_io);
    struct client *c = w->data;

    if ((n = read(w->fd, c->req->buf.raw, c->req->buf.size)) < 0) {
        if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
            return;
        client_log_socket_error(LOG_WARNING, c, w->fd, "recv error");
        goto DESTROY;
    }
    if (!n)
        goto DESTROY;
    c->req->buf.len += n;

    switch (c->req->buf.raw[0]) {
    case SOCKS4:
        cb = socks4_addr_handler;
        break;
    case SOCKS5:
        cb = socks5_hello_handler;
        break;
    case 'C': /* CONNECT/COPY */
    case 'D': /* DELETE */
    case 'G': /* GET */
    case 'H': /* HEAD */
    case 'L': /* LOCK */
    case 'M': /* MKCOL/MOVE */
    case 'O': /* OPTIONS */
    case 'P': /* POST/PUT/PATCH/PROPFIND/PROPPATCH */
    case 'T': /* TRACE */
    case 'U': /* UNLOCK */
        cb = http_header_handler;
        break;
    default:
        LOGW("%s unsupported protocol", c->laddrinfo);
        goto DESTROY;
    }

    ev_set_cb(w, cb);
    EV_CB_INVOKE(w, revents);

    return;

DESTROY:
    client_destroy(c);
}

void client_update_activity(struct client *c) {
    ev_timer_again(c->ctx->evloop, &c->evtimer);
}

static void accept_handler(EV_P_ ev_io *w, int revents)
{
    struct sockaddr_storage address;
    socklen_t address_len = sizeof(address);
    struct client_ctx *clictx = w->data;

    int conn = accept(w->fd, (struct sockaddr *)&address, &address_len);
    if (conn < 0) {
        switch (errno) {
#if EAGAIN != EWOULDBLOCK
        case EWOULDBLOCK:
#endif /* #if EAGAIN != EWOULDBLOCK */
        case EAGAIN: case EINTR: case ECONNABORTED:
            return;
        case ETIMEDOUT: case EMFILE: case ENFILE: case ENOBUFS: case ENOMEM:
            LOGW("accept error: %s", strerror(errno));
            return;
        default:
            LOGE("accept error: %s", strerror(errno));
            exit(EXIT_FAILURE);
        }
    }

    struct client *c = create(clictx);
    if (c) {
        c->lsock = conn;
        sockaddr_ntop(&address, c->laddrinfo, sizeof(c->laddrinfo));
        nonblocking(conn);
        ev_io_init(&c->levio, hello_handler, conn, EV_READ);
        ev_io_start(EV_A_ &c->levio);
        ev_timer_init(&c->evtimer, sweep_handler, 0., clictx->cfg->accept_timeout);
        ev_timer_again(EV_A_ &c->evtimer);
    } else {
        LOGW("Not enough memory");
        close(conn);
    }
}

void sweep_handler(EV_P_ ev_timer *watcher, int revents)
{
    struct client *c = watcher->data;
    struct linger linger = { .l_onoff = 1, .l_linger = 0 };
    LOGI("Connection \"%s %s --> %s --> %s\" no data long time, close it.",
            src_proto_str(c->src_proto), c->laddrinfo, c->srv->addr, c->dst_info);
    setsockopt(c->lsock, SOL_SOCKET, SO_LINGER, &linger, sizeof(linger));
    client_destroy(c);
}

struct client_ctx *client_ctx_create(const struct cfg *cfg,
        struct resolv_ctx *resolver, struct ev_loop *evloop)
{
    struct client_ctx *cctx;

    if (!(cctx = malloc(sizeof(*cctx) + cfg->buf_size * 2))) {
        LOGE("Not enough memory");
        return NULL;
    }

    cctx->listen_sock = create_listen_socket(cfg->local_bind, cfg->local_listen,
            AF_UNSPEC, SOCK_STREAM);
    if (cctx->listen_sock < 0) {
        free(cctx);
        return NULL;
    }

    cctx->cfg = cfg;
    cctx->resolver = resolver;
    TAILQ_INIT(&cctx->client_pool.inuse);
    TAILQ_INIT(&cctx->client_pool.unuse);
    cctx->client_pool.unuseTotal = 0;
    cctx->rand_seed = cfg->rand_seed;
    cctx->inbuflen = cctx->outbuflen = cfg->buf_size;
    cctx->inbuf = (void *)cctx + sizeof(*cctx);
    cctx->outbuf = (void *)cctx + sizeof(*cctx) + cctx->inbuflen;

    cctx->evloop = evloop;
    ev_io_init(&cctx->evio, accept_handler, cctx->listen_sock, EV_READ);
    cctx->evio.data = cctx;
    ev_io_start(evloop, &cctx->evio);

    return cctx;
}

int client_ctx_destroy(struct client_ctx *cctx)
{
    struct client *c;

    while ((c = TAILQ_LAST(&cctx->client_pool.inuse, inuse)))
        client_destroy(c);

    while ((c = TAILQ_LAST(&cctx->client_pool.unuse, unuse)))
        client_free(c);

    if (cctx->listen_sock >= 0) {
        ev_io_stop(cctx->evloop, &cctx->evio);
        close(cctx->listen_sock);
    }

    free(cctx);

    return 0;
}
